#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
I3D模型实现
基于膨胀3D卷积网络的微表情识别模型
参考: https://arxiv.org/abs/1705.07750
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class Conv3DBlock(nn.Module):
    """3D卷积块"""
    
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
        super(Conv3DBlock, self).__init__()
        
        self.conv = nn.Conv3d(
            in_channels, out_channels,
            kernel_size=kernel_size, stride=stride, padding=padding
        )
        self.bn = nn.BatchNorm3d(out_channels)
        self.relu = nn.ReLU(inplace=True)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x


class InceptionBlock3D(nn.Module):
    """3D Inception块"""
    
    def __init__(self, in_channels, out_channels):
        super(InceptionBlock3D, self).__init__()
        
        # 分支1: 1x1x1
        self.branch1 = Conv3DBlock(in_channels, out_channels[0], kernel_size=1, padding=0)
        
        # 分支2: 1x1x1 -> 3x3x3
        self.branch2_1 = Conv3DBlock(in_channels, out_channels[1], kernel_size=1, padding=0)
        self.branch2_2 = Conv3DBlock(out_channels[1], out_channels[2], kernel_size=3, padding=1)
        
        # 分支3: 1x1x1 -> 3x3x3 -> 3x3x3
        self.branch3_1 = Conv3DBlock(in_channels, out_channels[3], kernel_size=1, padding=0)
        self.branch3_2 = Conv3DBlock(out_channels[3], out_channels[4], kernel_size=3, padding=1)
        self.branch3_3 = Conv3DBlock(out_channels[4], out_channels[5], kernel_size=3, padding=1)
        
        # 分支4: 3x3x3池化 -> 1x1x1
        self.branch4_1 = nn.MaxPool3d(kernel_size=3, stride=1, padding=1)
        self.branch4_2 = Conv3DBlock(in_channels, out_channels[6], kernel_size=1, padding=0)
    
    def forward(self, x):
        branch1 = self.branch1(x)
        
        branch2 = self.branch2_1(x)
        branch2 = self.branch2_2(branch2)
        
        branch3 = self.branch3_1(x)
        branch3 = self.branch3_2(branch3)
        branch3 = self.branch3_3(branch3)
        
        branch4 = self.branch4_1(x)
        branch4 = self.branch4_2(branch4)
        
        outputs = [branch1, branch2, branch3, branch4]
        return torch.cat(outputs, 1)


class I3D(nn.Module):
    """
    膨胀3D卷积网络 (Inflated 3D ConvNet)
    用于微表情视频序列分类
    """
    
    def __init__(self, num_classes=7, pretrained=False, dropout_prob=0.5):
        super(I3D, self).__init__()
        
        self.num_classes = num_classes
        
        # 初始卷积层
        self.conv1 = Conv3DBlock(3, 64, kernel_size=7, stride=2, padding=3)
        self.pool1 = nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1))
        
        # Conv2
        self.conv2 = Conv3DBlock(64, 64, kernel_size=1, padding=0)
        self.conv3 = Conv3DBlock(64, 192, kernel_size=3, padding=1)
        self.pool2 = nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1))
        
        # Inception模块
        # Inception3a
        self.inception3a = InceptionBlock3D(192, [64, 96, 128, 16, 32, 32, 32])
        # Inception3b
        self.inception3b = InceptionBlock3D(256, [128, 128, 192, 32, 96, 64, 64])
        # MaxPool
        self.pool3 = nn.MaxPool3d(kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1))
        
        # Inception4a-e
        self.inception4a = InceptionBlock3D(480, [192, 96, 208, 16, 48, 64, 64])
        self.inception4b = InceptionBlock3D(512, [160, 112, 224, 24, 64, 64, 64])
        self.inception4c = InceptionBlock3D(512, [128, 128, 256, 24, 64, 64, 64])
        self.inception4d = InceptionBlock3D(512, [112, 144, 288, 32, 64, 64, 64])
        self.inception4e = InceptionBlock3D(528, [256, 160, 320, 32, 128, 128, 128])
        # MaxPool
        self.pool4 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2), padding=(0, 0, 0))
        
        # Inception5a-b
        self.inception5a = InceptionBlock3D(832, [256, 160, 320, 32, 128, 128, 128])
        self.inception5b = InceptionBlock3D(832, [384, 192, 384, 48, 128, 128, 128])
        
        # 全局平均池化
        self.avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
        self.dropout = nn.Dropout(dropout_prob)
        
        # 分类器
        self.fc = nn.Linear(1024, num_classes)
        
        # 权重初始化
        self._initialize_weights()
        
        # 加载预训练权重
        if pretrained:
            self._load_pretrained_weights()
    
    def forward(self, x):
        """
        前向传播
        
        Args:
            x: 输入张量 [batch_size, channels, frames, height, width]
        
        Returns:
            logits: 分类输出
        """
        # 确保输入是5D张量 [B, C, T, H, W]
        if x.dim() == 4:
            x = x.unsqueeze(2)  # 添加时间维度
        
        # 初始层
        x = self.conv1(x)
        x = self.pool1(x)
        
        # Conv2
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.pool2(x)
        
        # Inception3
        x = self.inception3a(x)
        x = self.inception3b(x)
        x = self.pool3(x)
        
        # Inception4
        x = self.inception4a(x)
        x = self.inception4b(x)
        x = self.inception4c(x)
        x = self.inception4d(x)
        x = self.inception4e(x)
        x = self.pool4(x)
        
        # Inception5
        x = self.inception5a(x)
        x = self.inception5b(x)
        
        # 全局平均池化
        x = self.avg_pool(x)
        x = x.view(x.size(0), -1)  # 展平
        
        # Dropout
        x = self.dropout(x)
        
        # 分类器
        x = self.fc(x)
        
        return x
    
    def _initialize_weights(self):
        """初始化网络权重"""
        for m in self.modules():
            if isinstance(m, nn.Conv3d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm3d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)
    
    def _load_pretrained_weights(self):
        """
        加载预训练权重
        这里只是一个框架，实际使用时需要根据具体的预训练模型进行实现
        """
        # 加载预训练权重
        try:
            # 这里可以添加加载预训练I3D权重的代码
            # 例如从Kinetics数据集上预训练的I3D模型
            # pretrained_dict = torch.load('path/to/pretrained/i3d.pth')
            # model_dict = self.state_dict()
            # ...
            pass
        except Exception as e:
            print(f"加载预训练权重失败: {e}")
            print("使用随机初始化的权重")
            pass 